#include "filter/filter.h"

#include <stdlib.h>
#include <math.h>

#include "SLAM/Probabilities.h"
#include <visualization_msgs/Marker.h>

Filter::Filter(int num_particles) :
	robot_dx_(0),
	robot_dy_(0),
	robot_dtheta_(0),
	num_particles_(num_particles),
	particles_(num_particles),
	cumulative_sum_(num_particles)
{
	// Initialize particles
	for (int i = 0; i < num_particles; i++) {
		particles_[i].x = particles_[i].y = particles_[i].theta = 0.0f;
	}

	// Create publisher to publish visualization calls
	ros::NodeHandle nh;
	particle_publisher_ = nh.advertise<visualization_msgs::Marker>("/particles", 0);
	probabilities_client_ = nh.serviceClient<SLAM::Probabilities>("/slam/particle_probabilities");
	phase_subscriber_ = nh.subscribe("/SLAM/phase", 100, &Filter::switch_phase_callback, this);
}

Filter::~Filter()
{

}

void Filter::switch_phase_callback(const SLAM::SwitchPhase::ConstPtr& msg) {
	if (msg->phase == 1 || msg->phase == 3) {
		for (int i = 0; i < num_particles_; i++) {
			particles_[i].x = particles_[i].y = particles_[i].theta = 0.0f;
		}
		robot_dx_ = robot_dy_ = robot_dtheta_ = 0.0f;
		estimate_robot_x_ = estimate_robot_y_ = estimate_robot_theta_ = 0.0f;
	}
}

float gaussrand()
{
	static float V1, V2, S;
	static int phase = 0;
	float X;

	if (phase == 0) {
		do {
			float U1 = (float) rand() / RAND_MAX;
			float U2 = (float) rand() / RAND_MAX;

			V1 = 2 * U1 - 1;
			V2 = 2 * U2 - 1;
			S = V1 * V1 + V2 * V2;
		} while(S >= 1 || S == 0);

		X = V1 * sqrtf(-2 * log(S) / S);
	} else
		X = V2 * sqrtf(-2 * log(S) / S);

	phase = 1 - phase;

	return X;
}

float normal_dist(float mean, float std_dev)
{
	return std_dev * gaussrand() + mean;
}

//void Filter::robotMoved(float delta_x, float delta_y, float delta_theta)
//{
//	// Step the particles with uncertainty in moved distance
//	for (int i = 0; i < num_particles_; i++) {
//		// Temp testing
//		// Since delta_x and delta_y are in absolute coordinates these calcuations below
//		// are quite faulty (they don't consider the orientation of each particle)
//		particles_[i].x += normal_dist(delta_x, 0.0025f);
//		particles_[i].y += normal_dist(delta_y, 0.0025f);
//		particles_[i].theta += normal_dist(delta_theta, 0.01f);
//	}
//}

void Filter::robotMoved(float dist_left, float dist_right, float wheel_base)
{
	// Step the particles with uncertainty in moved distance
	for (int i = 0; i < num_particles_; i++) {
		float uncertain_dist_left = normal_dist(dist_left, 0.005f);
		float uncertain_dist_right = normal_dist(dist_right, 0.005f);

		float dist_center = (uncertain_dist_left + uncertain_dist_right) / 2.0f;
		float delta_theta = -(uncertain_dist_right - uncertain_dist_left) / wheel_base;
		float delta_x = dist_center * cosf(particles_[i].theta + delta_theta);
		float delta_y = dist_center * sinf(particles_[i].theta + delta_theta);

		particles_[i].x += delta_x;
		particles_[i].y += delta_y;
		particles_[i].theta += delta_theta;
	}
}

void Filter::estimatedPosition(float *x, float *y, float *theta)
{
	resampleParticles();

	*x = estimate_robot_x_;
	*y = estimate_robot_y_;
	*theta = estimate_robot_theta_;
}

void Filter::cachedEstimatedPosition(float *x, float *y, float *theta)
{
	*x = estimate_robot_x_;
	*y = estimate_robot_y_;
	*theta = estimate_robot_theta_;
}

void Filter::resampleParticles()
{
	// Load the particles into the particle probability service request
	SLAM::Probabilities probabilitiesSrv;
	std::vector<float> sent_particles(3 * num_particles_); // 3 floats per particle
	for (int i = 0; i < num_particles_; i++) {
		sent_particles[(i * 3)] = particles_[i].x;
		sent_particles[(i * 3) + 1] = particles_[i].y;
		sent_particles[(i * 3) + 2] = particles_[i].theta;
	}
	probabilitiesSrv.request.particles = sent_particles;

	// Probabilites for each particle now received from server, resample particles
	if (probabilities_client_.call(probabilitiesSrv)) {
		if (probabilitiesSrv.response.probabilities.size() == 0) return;

		// Calculate cumulative sum
		cumulative_sum_[0] = probabilitiesSrv.response.probabilities[0];
		for (int i = 1; i < num_particles_; i++) {
			uint8_t prob = probabilitiesSrv.response.probabilities[i];
			cumulative_sum_[i] = cumulative_sum_[i - 1] + prob;
		}

		// Find weighted mean of particles
		float x = 0;
		float y = 0;
		float theta = 0;
		float prob_sum = 0;
		for (int i = 0; i < num_particles_; i++) {
			float prob = (probabilitiesSrv.response.probabilities[i] / 255.0f);
			x += particles_[i].x * prob;
			y += particles_[i].y * prob;
			theta += particles_[i].theta * prob;
			prob_sum += prob;
		}
		estimate_robot_x_ = x / prob_sum;
		estimate_robot_y_ = y / prob_sum;
		estimate_robot_theta_ = theta / prob_sum;

		// Resample particles
		std::vector<int>::iterator sample_it;
		for (int i = 0; i < num_particles_; i++) {
			// Sample values between 0 and sum of probabilities
			int samp = cumulative_sum_[num_particles_ - 1] * ((float) rand() / RAND_MAX);
			sample_it = std::lower_bound(cumulative_sum_.begin(), cumulative_sum_.end(), samp);

			int index = sample_it - cumulative_sum_.begin();
			particles_[i].x = sent_particles[(index * 3)];
			particles_[i].y = sent_particles[(index * 3) + 1];
			particles_[i].theta = sent_particles[(index * 3) + 2];
		}
	} else {
		ROS_ERROR("filter: Could not communicate with /slam/particle_probabilities service.");
	}

	renderParticles();
}

void Filter::renderParticles()
{
	visualization_msgs::Marker marker;

	marker.header.frame_id = "/floor";
	marker.header.stamp = ros::Time();

	marker.ns = "particles";
	marker.id = 0;
	marker.lifetime = ros::Duration(0.25f);

	marker.type = visualization_msgs::Marker::LINE_LIST;
	marker.action = visualization_msgs::Marker::ADD;

	marker.pose.position.x = 0.0f;
	marker.pose.position.y = 0.0f;
	marker.pose.position.z = 0.0f;

	marker.pose.orientation.x = 0.0f;
	marker.pose.orientation.y = 0.0f;
	marker.pose.orientation.z = 0.0f;
	marker.pose.orientation.w = 1.0f;

	marker.scale.x = 0.01f;
	marker.scale.y = 1.0f;
	marker.scale.z = 1.0f;

	marker.color.r = 1.0f;
	marker.color.g = 0.5f;
	marker.color.b = 0.0f;
	marker.color.a = 1.0f;

	std::vector<geometry_msgs::Point> points(num_particles_ * 2);
	for (int i = 0; i < num_particles_; i += 2) {
		points[i].x = particles_[i].x;
		points[i].y = particles_[i].y;
		points[i].z = 0.0f;
		points[i + 1].x = particles_[i].x + 0.05f * cosf(particles_[i].theta);
		points[i + 1].y = particles_[i].y + 0.05f * sinf(particles_[i].theta);
		points[i + 1].z = 0.0f;
	}

	marker.points = points;

	// Publish the ray
	particle_publisher_.publish(marker);
}
